#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Advanced Joycon Controller Reader
使用专门的Joycon库读取手柄的陀螺仪、加速度计等传感器数据

Author: Yinchuan Li
Date: 2025-07-30
Description: 使用joycon-python库读取Joycon手柄的完整传感器数据
"""

import time
import sys
import math
import threading
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
import numpy as np

try:
    from joycon import JoyCon, get_L_id, get_R_id
    from joycon import ButtonEventJoyCon, GyroTrackingJoyCon
    JOYCON_AVAILABLE = True
except ImportError:
    print("警告: joycon-python库未安装，将使用pygame作为备选方案")
    print("安装命令: pip install joycon-python")
    JOYCON_AVAILABLE = False
    import pygame

@dataclass
class GyroscopeData:
    """陀螺仪数据结构"""
    x: float  # 绕X轴角速度 (度/秒)
    y: float  # 绕Y轴角速度 (度/秒)
    z: float  # 绕Z轴角速度 (度/秒)
    magnitude: float  # 角速度大小
    timestamp: float  # 时间戳

@dataclass
class AccelerometerData:
    """加速度计数据结构"""
    x: float  # X轴加速度 (m/s²)
    y: float  # Y轴加速度 (m/s²)
    z: float  # Z轴加速度 (m/s²)
    magnitude: float  # 加速度大小
    pitch: float  # 俯仰角 (度)
    roll: float  # 横滚角 (度)
    timestamp: float  # 时间戳

@dataclass
class JoyconState:
    """Joycon完整状态"""
    buttons: Dict[str, bool]
    sticks: Dict[str, Tuple[float, float]]
    gyroscope: GyroscopeData
    accelerometer: AccelerometerData
    battery_level: float
    timestamp: float

class AdvancedJoyconReader:
    """高级Joycon读取器类"""
    
    def __init__(self, joycon_type: str = "auto"):
        """
        初始化高级Joycon读取器
        
        Args:
            joycon_type: "left", "right", "auto" 或 "both"
        """
        self.joycon_type = joycon_type
        self.left_joycon = None
        self.right_joycon = None
        self.use_pygame = not JOYCON_AVAILABLE
        
        if self.use_pygame:
            self._init_pygame()
        else:
            self._init_joycon_library()
        
        # 数据存储
        self.current_state = None
        self.data_lock = threading.Lock()
        
        # 校准数据
        self.gyro_calibration = {'x': 0.0, 'y': 0.0, 'z': 0.0}
        self.accel_calibration = {'x': 0.0, 'y': 0.0, 'z': 0.0}
        
        # 数据平滑
        self.gyro_history = []
        self.accel_history = []
        self.max_history = 10
        
    def _init_pygame(self):
        """初始化pygame作为备选方案"""
        pygame.init()
        pygame.joystick.init()
        
        if pygame.joystick.get_count() == 0:
            print("错误: 未检测到手柄设备")
            sys.exit(1)
        
        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()
        print(f"使用pygame模式连接手柄: {self.joystick.get_name()}")
    
    def _init_joycon_library(self):
        """初始化专门的Joycon库"""
        try:
            if self.joycon_type in ["left", "auto", "both"]:
                left_id = get_L_id()
                if left_id:
                    self.left_joycon = GyroTrackingJoyCon(*left_id)
                    print(f"左Joycon连接成功: {left_id}")
            
            if self.joycon_type in ["right", "auto", "both"]:
                right_id = get_R_id()
                if right_id:
                    self.right_joycon = GyroTrackingJoyCon(*right_id)
                    print(f"右Joycon连接成功: {right_id}")
            
            if not self.left_joycon and not self.right_joycon:
                print("错误: 未检测到Joycon设备")
                print("请确保Joycon已通过蓝牙连接")
                sys.exit(1)
                
        except Exception as e:
            print(f"Joycon库初始化失败: {e}")
            print("切换到pygame模式")
            self.use_pygame = True
            self._init_pygame()
    
    def calibrate_sensors(self, duration: float = 5.0):
        """校准传感器"""
        print(f"开始传感器校准，请保持手柄静止 {duration} 秒...")
        
        gyro_samples = []
        accel_samples = []
        
        start_time = time.time()
        while time.time() - start_time < duration:
            if self.use_pygame:
                # pygame模式下的校准
                for event in pygame.event.get():
                    pass
                
                if self.joystick.get_numaxes() > 8:
                    gyro_x = self.joystick.get_axis(6)
                    gyro_y = self.joystick.get_axis(7)
                    gyro_z = self.joystick.get_axis(8)
                    gyro_samples.append((gyro_x, gyro_y, gyro_z))
                
                if self.joystick.get_numaxes() > 11:
                    accel_x = self.joystick.get_axis(9)
                    accel_y = self.joystick.get_axis(10)
                    accel_z = self.joystick.get_axis(11)
                    accel_samples.append((accel_x, accel_y, accel_z))
            else:
                # Joycon库模式下的校准
                if self.left_joycon:
                    gyro_data = self.left_joycon.get_gyro_data()
                    accel_data = self.left_joycon.get_accel_data()
                    if gyro_data:
                        gyro_samples.append((gyro_data['x'], gyro_data['y'], gyro_data['z']))
                    if accel_data:
                        accel_samples.append((accel_data['x'], accel_data['y'], accel_data['z']))
            
            time.sleep(0.1)
        
        # 计算校准值
        if gyro_samples:
            avg_gyro = np.mean(gyro_samples, axis=0)
            self.gyro_calibration = {'x': avg_gyro[0], 'y': avg_gyro[1], 'z': avg_gyro[2]}
            print(f"陀螺仪校准完成: X={avg_gyro[0]:.4f}, Y={avg_gyro[1]:.4f}, Z={avg_gyro[2]:.4f}")
        
        if accel_samples:
            avg_accel = np.mean(accel_samples, axis=0)
            self.accel_calibration = {'x': avg_accel[0], 'y': avg_accel[1], 'z': avg_accel[2]}
            print(f"加速度计校准完成: X={avg_accel[0]:.4f}, Y={avg_accel[1]:.4f}, Z={avg_accel[2]:.4f}")
    
    def _smooth_data(self, new_data: Tuple[float, float, float], history: List) -> Tuple[float, float, float]:
        """平滑传感器数据"""
        history.append(new_data)
        if len(history) > self.max_history:
            history.pop(0)
        
        if len(history) < 3:
            return new_data
        
        # 使用移动平均平滑数据
        smoothed = np.mean(history, axis=0)
        return tuple(smoothed)
    
    def read_joycon_data(self) -> JoyconState:
        """读取Joycon数据"""
        timestamp = time.time()
        
        if self.use_pygame:
            return self._read_pygame_data(timestamp)
        else:
            return self._read_joycon_library_data(timestamp)
    
    def _read_pygame_data(self, timestamp: float) -> JoyconState:
        """使用pygame读取数据"""
        # 处理事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return None
        
        # 读取按钮状态
        buttons = {}
        for i in range(self.joystick.get_numbuttons()):
            buttons[f"button_{i}"] = bool(self.joystick.get_button(i))
        
        # 读取摇杆状态
        sticks = {}
        if self.joystick.get_numaxes() >= 4:
            sticks['left'] = (self.joystick.get_axis(0), self.joystick.get_axis(1))
            sticks['right'] = (self.joystick.get_axis(2), self.joystick.get_axis(3))
        
        # 读取陀螺仪数据
        gyro_x = gyro_y = gyro_z = 0.0
        if self.joystick.get_numaxes() > 8:
            gyro_x = self.joystick.get_axis(6) - self.gyro_calibration['x']
            gyro_y = self.joystick.get_axis(7) - self.gyro_calibration['y']
            gyro_z = self.joystick.get_axis(8) - self.gyro_calibration['z']
        
        # 平滑陀螺仪数据
        gyro_data = self._smooth_data((gyro_x, gyro_y, gyro_z), self.gyro_history)
        gyro_magnitude = math.sqrt(sum(x*x for x in gyro_data))
        
        gyroscope = GyroscopeData(
            x=gyro_data[0], y=gyro_data[1], z=gyro_data[2],
            magnitude=gyro_magnitude, timestamp=timestamp
        )
        
        # 读取加速度计数据
        accel_x = accel_y = accel_z = 0.0
        if self.joystick.get_numaxes() > 11:
            accel_x = self.joystick.get_axis(9) - self.accel_calibration['x']
            accel_y = self.joystick.get_axis(10) - self.accel_calibration['y']
            accel_z = self.joystick.get_axis(11) - self.accel_calibration['z']
        
        # 平滑加速度计数据
        accel_data = self._smooth_data((accel_x, accel_y, accel_z), self.accel_history)
        accel_magnitude = math.sqrt(sum(x*x for x in accel_data))
        
        # 计算角度
        pitch = math.atan2(accel_data[1], math.sqrt(accel_data[0]**2 + accel_data[2]**2)) * 180 / math.pi
        roll = math.atan2(-accel_data[0], accel_data[2]) * 180 / math.pi
        
        accelerometer = AccelerometerData(
            x=accel_data[0], y=accel_data[1], z=accel_data[2],
            magnitude=accel_magnitude, pitch=pitch, roll=roll, timestamp=timestamp
        )
        
        return JoyconState(
            buttons=buttons,
            sticks=sticks,
            gyroscope=gyroscope,
            accelerometer=accelerometer,
            battery_level=0.0,  # pygame无法获取电池信息
            timestamp=timestamp
        )
    
    def _read_joycon_library_data(self, timestamp: float) -> JoyconState:
        """使用Joycon库读取数据"""
        buttons = {}
        sticks = {}
        gyro_data = None
        accel_data = None
        battery_level = 0.0
        
        # 读取左Joycon数据
        if self.left_joycon:
            try:
                # 按钮状态
                left_buttons = self.left_joycon.get_button_status()
                for btn, pressed in left_buttons.items():
                    buttons[f"left_{btn}"] = pressed
                
                # 摇杆状态
                left_stick = self.left_joycon.get_stick_status()
                if left_stick:
                    sticks['left'] = (left_stick['horizontal'], left_stick['vertical'])
                
                # 陀螺仪数据
                gyro = self.left_joycon.get_gyro_data()
                if gyro:
                    gyro_x = gyro['x'] - self.gyro_calibration['x']
                    gyro_y = gyro['y'] - self.gyro_calibration['y']
                    gyro_z = gyro['z'] - self.gyro_calibration['z']
                    
                    # 平滑数据
                    smoothed_gyro = self._smooth_data((gyro_x, gyro_y, gyro_z), self.gyro_history)
                    gyro_magnitude = math.sqrt(sum(x*x for x in smoothed_gyro))
                    
                    gyro_data = GyroscopeData(
                        x=smoothed_gyro[0], y=smoothed_gyro[1], z=smoothed_gyro[2],
                        magnitude=gyro_magnitude, timestamp=timestamp
                    )
                
                # 加速度计数据
                accel = self.left_joycon.get_accel_data()
                if accel:
                    accel_x = accel['x'] - self.accel_calibration['x']
                    accel_y = accel['y'] - self.accel_calibration['y']
                    accel_z = accel['z'] - self.accel_calibration['z']
                    
                    # 平滑数据
                    smoothed_accel = self._smooth_data((accel_x, accel_y, accel_z), self.accel_history)
                    accel_magnitude = math.sqrt(sum(x*x for x in smoothed_accel))
                    
                    # 计算角度
                    pitch = math.atan2(smoothed_accel[1], math.sqrt(smoothed_accel[0]**2 + smoothed_accel[2]**2)) * 180 / math.pi
                    roll = math.atan2(-smoothed_accel[0], smoothed_accel[2]) * 180 / math.pi
                    
                    accel_data = AccelerometerData(
                        x=smoothed_accel[0], y=smoothed_accel[1], z=smoothed_accel[2],
                        magnitude=accel_magnitude, pitch=pitch, roll=roll, timestamp=timestamp
                    )
                
                # 电池信息
                battery = self.left_joycon.get_battery_status()
                if battery:
                    battery_level = battery.get('level', 0.0)
                    
            except Exception as e:
                print(f"读取左Joycon数据失败: {e}")
        
        # 读取右Joycon数据
        if self.right_joycon:
            try:
                # 按钮状态
                right_buttons = self.right_joycon.get_button_status()
                for btn, pressed in right_buttons.items():
                    buttons[f"right_{btn}"] = pressed
                
                # 摇杆状态
                right_stick = self.right_joycon.get_stick_status()
                if right_stick:
                    sticks['right'] = (right_stick['horizontal'], right_stick['vertical'])
                
                # 如果左Joycon没有陀螺仪数据，使用右Joycon的
                if not gyro_data:
                    gyro = self.right_joycon.get_gyro_data()
                    if gyro:
                        gyro_x = gyro['x'] - self.gyro_calibration['x']
                        gyro_y = gyro['y'] - self.gyro_calibration['y']
                        gyro_z = gyro['z'] - self.gyro_calibration['z']
                        
                        smoothed_gyro = self._smooth_data((gyro_x, gyro_y, gyro_z), self.gyro_history)
                        gyro_magnitude = math.sqrt(sum(x*x for x in smoothed_gyro))
                        
                        gyro_data = GyroscopeData(
                            x=smoothed_gyro[0], y=smoothed_gyro[1], z=smoothed_gyro[2],
                            magnitude=gyro_magnitude, timestamp=timestamp
                        )
                
                # 如果左Joycon没有加速度计数据，使用右Joycon的
                if not accel_data:
                    accel = self.right_joycon.get_accel_data()
                    if accel:
                        accel_x = accel['x'] - self.accel_calibration['x']
                        accel_y = accel['y'] - self.accel_calibration['y']
                        accel_z = accel['z'] - self.accel_calibration['z']
                        
                        smoothed_accel = self._smooth_data((accel_x, accel_y, accel_z), self.accel_history)
                        accel_magnitude = math.sqrt(sum(x*x for x in smoothed_accel))
                        
                        pitch = math.atan2(smoothed_accel[1], math.sqrt(smoothed_accel[0]**2 + smoothed_accel[2]**2)) * 180 / math.pi
                        roll = math.atan2(-smoothed_accel[0], smoothed_accel[2]) * 180 / math.pi
                        
                        accel_data = AccelerometerData(
                            x=smoothed_accel[0], y=smoothed_accel[1], z=smoothed_accel[2],
                            magnitude=accel_magnitude, pitch=pitch, roll=roll, timestamp=timestamp
                        )
                        
            except Exception as e:
                print(f"读取右Joycon数据失败: {e}")
        
        # 如果没有获取到传感器数据，创建默认值
        if not gyro_data:
            gyro_data = GyroscopeData(0.0, 0.0, 0.0, 0.0, timestamp)
        if not accel_data:
            accel_data = AccelerometerData(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, timestamp)
        
        return JoyconState(
            buttons=buttons,
            sticks=sticks,
            gyroscope=gyro_data,
            accelerometer=accel_data,
            battery_level=battery_level,
            timestamp=timestamp
        )
    
    def print_status(self, state: JoyconState):
        """打印当前状态"""
        if state is None:
            return False
        
        # 清屏
        print("\033[2J\033[H")
        
        print("=" * 80)
        print("Advanced Joycon Controller Reader")
        print("Author: Assistant | Date: 2024-12-19")
        print("=" * 80)
        print("按 Ctrl+C 退出程序 | 按 'c' 校准传感器")
        print("-" * 80)
        
        # 显示按钮状态
        print("🔘 按钮状态:")
        print("-" * 40)
        active_buttons = [btn for btn, pressed in state.buttons.items() if pressed]
        if active_buttons:
            print("活动按钮:", ", ".join(active_buttons))
        else:
            print("无按钮按下")
        print()
        
        # 显示摇杆状态
        print("🎮 摇杆状态:")
        print("-" * 40)
        for stick_name, (x, y) in state.sticks.items():
            print(f"{stick_name.capitalize()} Stick: X={x:6.3f}, Y={y:6.3f}")
        print()
        
        # 显示陀螺仪数据
        print("🔄 陀螺仪数据:")
        print("-" * 40)
        gyro = state.gyroscope
        print(f"X轴角速度: {gyro.x:8.3f} °/s")
        print(f"Y轴角速度: {gyro.y:8.3f} °/s")
        print(f"Z轴角速度: {gyro.z:8.3f} °/s")
        print(f"角速度大小: {gyro.magnitude:8.3f} °/s")
        print()
        
        # 显示加速度计数据
        print("📊 加速度计数据:")
        print("-" * 40)
        accel = state.accelerometer
        print(f"X轴加速度: {accel.x:8.3f} m/s²")
        print(f"Y轴加速度: {accel.y:8.3f} m/s²")
        print(f"Z轴加速度: {accel.z:8.3f} m/s²")
        print(f"加速度大小: {accel.magnitude:8.3f} m/s²")
        print(f"俯仰角: {accel.pitch:8.1f}°")
        print(f"横滚角: {accel.roll:8.1f}°")
        print()
        
        # 显示电池状态
        if state.battery_level > 0:
            print("🔋 电池状态:")
            print("-" * 40)
            print(f"电池电量: {state.battery_level:.1f}%")
            print()
        
        # 显示活动摘要
        print("📊 活动摘要:")
        print("-" * 40)
        activities = []
        
        # 检查摇杆活动
        for stick_name, (x, y) in state.sticks.items():
            if abs(x) > 0.1 or abs(y) > 0.1:
                activities.append(f"{stick_name}摇杆活动")
        
        # 检查陀螺仪活动
        if abs(gyro.magnitude) > 5.0:
            activities.append(f"陀螺仪活动: {gyro.magnitude:.1f}°/s")
        
        # 检查加速度计活动
        if abs(accel.magnitude - 9.81) > 2.0:
            activities.append(f"加速度变化: {accel.magnitude:.1f}m/s²")
        
        if active_buttons:
            activities.extend(active_buttons)
        
        if activities:
            print("活动:", ", ".join(activities))
        else:
            print("无活动")
        
        print("=" * 80)
        return True
    
    def run(self, update_rate: float = 30.0):
        """运行主循环"""
        print("启动高级Joycon读取器...")
        print("按 Ctrl+C 退出程序")
        print("按 'c' 校准传感器")
        time.sleep(1)
        
        try:
            while True:
                state = self.read_joycon_data()
                if not self.print_status(state):
                    break
                time.sleep(1.0 / update_rate)
                
        except KeyboardInterrupt:
            print("\n程序已退出")
        finally:
            if self.use_pygame:
                pygame.quit()
            else:
                if self.left_joycon:
                    self.left_joycon.close()
                if self.right_joycon:
                    self.right_joycon.close()


def main():
    """主函数"""
    try:
        # 创建高级Joycon读取器
        reader = AdvancedJoyconReader()
        
        # 询问是否校准传感器
        print("是否要校准传感器？(y/n): ", end="")
        try:
            import msvcrt  # Windows
            if msvcrt.getch().decode().lower() == 'y':
                reader.calibrate_sensors()
        except ImportError:
            try:
                import tty
                import termios
                
                # Linux/Mac
                fd = sys.stdin.fileno()
                old_settings = termios.tcgetattr(fd)
                try:
                    tty.setraw(sys.stdin.fileno())
                    ch = sys.stdin.read(1)
                    if ch.lower() == 'y':
                        reader.calibrate_sensors()
                finally:
                    termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
            except:
                pass
        
        # 运行主循环
        reader.run()
        
    except Exception as e:
        print(f"错误: {e}")
        print("请确保:")
        print("1. Joycon手柄已正确连接到电脑")
        print("2. 如果使用蓝牙连接，确保手柄已配对")
        print("3. 安装joycon-python库以获得更好的支持: pip install joycon-python")


if __name__ == "__main__":
    main() 